import os, sys
print(os.getcwd())
sys.path.append(os.getcwd())

from skorch.core import *
from torch.optim import SGD, Adam
from torch.nn.functional import binary_cross_entropy_with_logits

from skorch.vision import *
from skorch import NeuralNet
from skorch.helper import predefined_split
from skorch.callbacks import EpochScoring, LRScheduler, CyclicLR, Checkpoint
from skorch.loss import *

from model import unet
from dataset import create_dr_segment_datasets, create_segment_multi_augmenter, jsrt_stats, denormalize_image
torch.manual_seed(0)


def approximate_iou_metric(true_masks:NPArrayMask, predicted_masks:NPArrayMask, sigmoid=True):
    "Jaccard similarity coefficient score"
    if sigmoid:
        test_thresh = 0.5
        predicted_masks = 1 / (1 + np.exp(-predicted_masks))
        predicted_masks = (predicted_masks > test_thresh).astype(np.uint8)

    approx_intersect = np.sum(np.minimum(predicted_masks, true_masks), axis=(2, 3))
    approx_union = np.sum(np.maximum(predicted_masks, true_masks), axis=(2, 3))
    return np.around(np.mean(approx_intersect / approx_union, axis=0), decimals=4)


def approximate_dice_metric(true_masks:NPArrayMask, predicted_masks:NPArrayMask, sigmoid=True):
    "Dice metric"
    "https://github.com/mattmacy/torchbiomed/blob/master/torchbiomed/loss.py"
    if sigmoid:
        test_thresh = 0.5
        predicted_masks = 1 / (1 + np.exp(-predicted_masks))

        predicted_masks = (predicted_masks > test_thresh).astype(np.uint8)

    predicted_masks = predicted_masks.reshape((predicted_masks.shape[0], predicted_masks.shape[1], -1))
    true_masks = true_masks.reshape((true_masks.shape[0], true_masks.shape[1], -1)).astype(np.uint8)

    approx_intersect = (predicted_masks * true_masks)
    score = 2. * approx_intersect.sum(axis=2) / (predicted_masks.sum(axis=2) + true_masks.sum(axis=2))
    return np.around(score.mean(axis=0), decimals=4)


def iou_scoring(net, ds, y):
    predicted_logit_masks = net.predict(ds)
    return approximate_iou_metric(y, predicted_logit_masks)


def dice_scoring(net, ds, y):
    predicted_logit_masks = net.predict(ds)
    return approximate_dice_metric(y, predicted_logit_masks)

def get_dataloader():
    target_size = (615, 512)
    data_root = '/home/blake/data/dataset/datasets/DR/dicom/DR2'
    return create_dr_segment_datasets(data_root, create_segment_multi_augmenter, jsrt_stats,
                                                     target_size=target_size)

def test_model():
    #dr image h/w = 1.2
    target_size = (615, 512) #(h,w)
    pretrianed = True
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    is_parallel = True
    module = unet(n_classes=3, feature_scale=1, imsize=target_size, ce_encoder=True,debug=True)
    x = torch.ones((1, 3, 615, 512))
    y = module.forward(x)


def init_model():
    target_size = (615, 512)
    pretrianed = False
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    is_parallel = True
    module = unet(n_classes=3, feature_scale=1, imsize=target_size, ce_encoder=True)

    iou = EpochScoring(iou_scoring, name='iou', lower_is_better=False)
    dice = EpochScoring(dice_scoring, name='dice', lower_is_better=False)

    cyclicLR = LRScheduler(policy=CyclicLR,
                           base_lr=0.002,
                           max_lr=0.2,
                           step_size_up=540,
                           step_size_down=540)

    net = NeuralNet(
        module,
        criterion=BCEDiceLoss,
        batch_size=2,
        max_epochs=2,
        optimizer=Adam,
        lr=0.001,
        iterator_train__shuffle=True,
        iterator_train__num_workers=8,
        iterator_valid__shuffle=False,
        iterator_valid__num_workers=8,
        callbacks=[cyclicLR,
                   dice,
                   Checkpoint(f_params='best_params.pt', dirname='segment')],
        device=device,
        parallel=is_parallel
    )

    if pretrianed:
        net.initialize()
        net.load_params(f_params='unet_server/best_params.pt', f_optimizer='unet_server/optimizer.pt')

    return net


def train():
    train_ds, valida_ds = get_dataloader()
    net = init_model()
    net.fit(train_ds)


def predict():
    train_ds, valida_ds = get_dataloader()
    net = init_model()
    net.initialize()

    net.load_params(f_params='unet_server/best_params.pt')

    scores = []
    for i in range(len(valida_ds)):
        image, mask = valida_ds[i]
        input_image = image.unsqueeze_(0)
        input_mask = mask.unsqueeze_(0)
        predicted_mask = net.predict_tensor(input_image)
        predicted_mask = 1 / (1 + np.exp(-predicted_mask))
        predicted_mask = (predicted_mask > 0.5).astype(np.uint8)
        cur_score = approximate_iou_metric(to_np(input_mask), predicted_mask, sigmoid=False)
        scores.append(cur_score)

    scores = np.array(scores)

    #dice score  [0.95558526 0.92767437 0.91658807]
    #iou score  [0.91875346 0.86917337 0.85155573]

    print("score ",(scores.mean(axis=0)))


def predict_show():
    train_ds, valida_ds = get_dataloader()
    net = init_model()
    net.initialize()

    net.load_params(f_params='unet_server/best_params.pt')

    scores = []
    for i in range(len(valida_ds)):
        image, mask = valida_ds[i]
        input_image = image.unsqueeze_(0)
        input_mask = mask.unsqueeze_(0)
        predicted_mask = net.predict_tensor(input_image)
        predicted_mask = 1 / (1 + np.exp(-predicted_mask))
        predicted_mask = (predicted_mask > 0.1).astype(np.uint8) * 255
        cur_score = approximate_dice_metric(to_np(input_mask), predicted_mask, sigmoid=False)
        scores.append(cur_score)

    scores = np.array(scores)

    print("score ",(scores.mean(axis=0)))


def plot_mask_cells(mask_cells):
    """Plots cells with their true mask, predicted mask.

    Parameters
    ----------
    mask_cells: list of tuples (`true_mask`, `predicted_mask`, `cell`)
    padding: int (default=16)
        Padding around mask to remove.
    """
    fig, axes = plt.subplots(len(mask_cells), 7, figsize=(12, 10))
    for idx, (axes, mask_cell) in enumerate(zip(axes, mask_cells), 1):
        ax1, ax2, ax3, ax4, ax5, ax6, ax7 = axes
        true_mask, predicted_mask, cell = mask_cell
        plot_mask_cell(
                true_mask, predicted_mask, cell,
                'Type {}'.format(idx),
                ax1, ax2, ax3, ax4, ax5, ax6, ax7)
    fig.tight_layout()
    return fig, axes


def plot_mask_cell(true_mask,
                   predicted_mask,
                   cell,
                   suffix,
                   ax1, ax2, ax3, ax4, ax5, ax6, ax7):
    """Plots a single cell with a its true mask and predicuted mask"""
    for ax in [ax1, ax2, ax3, ax4, ax5, ax6, ax7]:
        ax.grid(False)
        ax.set_xticks([])
        ax.set_yticks([])

    show_image = denormalize_image(cell, jsrt_stats[0], jsrt_stats[1])
    show_image = np.transpose(show_image, [1,2,0])
    true_mask = to_detach(true_mask, cpu=True)
    true_mask = to_np(true_mask)

    ax1.imshow(show_image)
    ax1.set_title('Image - {}'.format(suffix))

    ax2.imshow(np.squeeze(true_mask[0]), cmap='viridis')
    ax2.set_title('True Mask Lung - {}'.format(suffix))
    ax3.imshow(np.squeeze(true_mask[1]), cmap='viridis')
    ax3.set_title('True Mask Clavicle - {}'.format(suffix))
    ax4.imshow(np.squeeze(true_mask[2]), cmap='viridis')
    ax4.set_title('True Mask Heart - {}'.format(suffix))

    ax5.imshow(predicted_mask[0], cmap='viridis')
    ax5.set_title('Predicted Lung - {}'.format(suffix))
    ax6.imshow(predicted_mask[1], cmap='viridis')
    ax6.set_title('Predicted Clavicle - {}'.format(suffix))
    ax7.imshow(predicted_mask[2], cmap='viridis')
    ax7.set_title('Predicted Heart - {}'.format(suffix))
    return ax1, ax2, ax3, ax4, ax5, ax6, ax7


def test_predict_image():
    net = init_model()
    net.initialize()

    net.load_params(f_params='segment/best_params.pt')

    image_path = '/home/blake/data/dataset/datasets/DR/dicom/DR/images/200-0000000001-0001-10001-1.3.51.0.7.11073525975.49120.64333.42989.22486.28979.41145.dcm'

    image = load_image_obj(str(image_path), 'pydcm')
    image = cv2.cvtColor(image, cv2.COLOR_GRAY2RGB)
    image = resize_image(image, (512, 512))

    show_image = np.copy(image)

    start_time = time.time()
    image = image_to_tensor(image, jsrt_stats)
    image = image.unsqueeze_(0)

    predicted_mask = net.predict(image)
    print('time: ', time.time() - start_time)
    predicted_mask = 1 / (1 + np.exp(-predicted_mask))
    predicted_mask = (predicted_mask*255).astype(np.uint8)
    predicted_mask = np.squeeze(predicted_mask)

    fig, axes = plt.subplots(1, 4, figsize=(50, 50))
    ax1, ax2, ax3, ax4 = axes

    ax1.imshow(show_image)
    ax1.set_title('Image')

    ax2.imshow(predicted_mask[0], cmap='gray')
    ax2.set_title('Predicted Lung')
    ax3.imshow(predicted_mask[1], cmap='gray')
    ax3.set_title('Predicted Clavicle')
    ax4.imshow(predicted_mask[2], cmap='gray')
    ax4.set_title('Predicted Heart')
    plt.savefig('show.png', bbox_inches='tight')
    plt.show()


if __name__ == '__main__':
    print('')

    #test_model()
    train()
    #test_predict_image()
    #predict()


